В этом параграфе мы снова попробуем решить задачу генерации, когда нам дана выборка объектов из распределения x0∼q(x), и хотим научиться генерировать новые объекты из распределения , которых нет в нашей выборке.
Вероятно, вы уже знакомы с другими генеративными моделями, например VAE или GAN-ы. Здесь же мы познакомим вас с еще одним видом генеративных моделей: диффузионные модели, которые стали крайне популярны в последнее время благодаря своему высокому качеству генерации объектов из заданного распределения. В общий чертах, они работают следующим образом: берем шум из N(0,I) и шаг за шагом удаляем компоненты шума до тех пор, пока не получим объект x0 из распределения, см. иллюстрацию ниже.
Более детально
Для детального понимания стоит объяснить, что такое прямой и обратный диффузионные процессы. Прямой процесс заключается в постепенном зашумлении картинки с помощью распределения q, а обратный, наоборот, в расшумлении с помощью распределения p. Их можно схематично изобразить следующим образом:
Прямой диффузионный процесс определяется как апостериорное распределение q(x1:T∣x0). Это распределение также является Марковской цепочкой, которая постепенно добавляет гауссовский шум к объекту x0. На каждом шаге шум добавляется с различной магнитудой, которая определяется расписанием дисперсий {β1,...,βT}. При правильном выборе расписания в пределе по числу шагов T мы должны сойтись к шуму из N(0,I). В качестве распределений q берут нормальные распределения:
Теперь перейдем к обратному процессу и к самой диффузионной модели.
Диффузионная модель - это вероятностная модель с латентными переменными вида pθ(x0):=∫pθ(x0:T)dx1:T, где промежуточные состояния x1,...,xT соответствуют зашумленным объектам, a x0 - объект из распределения. Совместное распределение pθ(x0:T) называет обратным диффузионным процессом, который представляет собой Марковскую цепочку из гауссовских распределений pθ(xi−1∣xi):
Таким образом, обратный процесс параметризуется моделью θ, которая по зашумленному объекту xt и шагу t предсказывает среднее μθ(xt,t) и дисперсию Σθ(xt,t).
Обучение диффузионной модели
Диффузионный модели обучаются, максимизируя вариационную нижнюю оценку (ELBO) логарифма правдоподобия logpθ(x0). По тому же принципу обучаются VAE, с тем лишь отличием, что у диффузионных моделей другая форма модели с латентными переменными. Итак, давайте выведем ELBO для диффузии:
Теперь вернемся к распределению q(xt∣xt−1). Для того чтобы получить xt, придется итеративно получать x1,...,xt−1. Однако это можно сделать более эффективно благодаря нормальным распределениям. Для этого обозначим αt:=1−βt и αˉt:=∏i=1tαi, тогда
Серым в скобках комментарий к последующему переходу.
LVLB=Eq(x0:T)[logpθ(x0:T)q(x1:T∣x0)](расписываемсовместноераспределение)=Eq[logpθ(xT)∏t=1Tpθ(xt−1∣xt)∏t=1Tq(xt∣xt−1)](беремлогарифм)=Eq[−logpθ(xT)+t=1∑Tlogpθ(xt−1∣xt)q(xt∣xt−1)](отщепляемчленысуммы)=Eq[−logpθ(xT)+t=2∑Tlogpθ(xt−1∣xt)q(xt∣xt−1)+logpθ(x0∣x1)q(x1∣x0)](*)=Eq[−logpθ(xT)+t=2∑Tlog(pθ(xt−1∣xt)q(xt−1∣xt,x0)⋅q(xt−1∣x0)q(xt∣x0))+logpθ(x0∣x1)q(x1∣x0)](логпроизведенияраскрываем)=Eq[−logpθ(xT)+t=2∑Tlogpθ(xt−1∣xt)q(xt−1∣xt,x0)+t=2∑Tlogq(xt−1∣x0)q(xt∣x0)+logpθ(x0∣x1)q(x1∣x0)](отвторойсуммыостанетсятолько 1ыйипоследнийчлен)=Eq[−logpθ(xT)+t=2∑Tlogpθ(xt−1∣xt)q(xt−1∣xt,x0)+logq(x1∣x0)q(xT∣x0)+logpθ(x0∣x1)q(x1∣x0)](комбинируем 1 и 3 член, 3 и 4 член)=Eq[logpθ(xT)q(xT∣x0)+t=2∑Tlogpθ(xt−1∣xt)q(xt−1∣xt,x0)−logpθ(x0∣x1)]=Eq[LTDKL(q(xT∣x0)∥pθ(xT))+t=2∑TLt−1DKL(q(xt−1∣xt,x0)∥pθ(xt−1∣xt))L0−logpθ(x0∣x1)]
Пояснение (*). Пользуемся тем, что у нас Марковский процесс, и теоремой Байеса:
Таким образом во время обучения, на каждой итерации параллельно оптимизируются случайные член Lt с помощью градиентного спуск (сэмлируем t∼U{1,...,T}). Поскольку все распределения нормальные, то KL между ними можно выписать в явной форме (см. ниже).
Осталось только выписать q(xt−1∣xt,x0) . Мы знаем, поскольку у нас все распределения нормальные, то и q(xt−1∣xt,x0) будет нормальным.
Обозначим
q(xt−1∣xt,x0)=N(xt−1;μ~(xt,x0),β~tI)
Вывод q(xt−1∣xt,x0)
Применим формулу Байеса и распишем. Тут мы просто пытаемся понять, как будут выглядеть среднее и дисперсия, выделяя квадратичную форму в показателе экспоненты
В прошлой подсекции наша модель предсказывала среднее и дисперсию нормального распределения. Давайте зафиксируем Σθ(xt,t)=σt2I. Обычно берут σt2=βt или σt2=β~t=1−αˉt1−αˉt−1βt. Тогда Lt−1 из предыдущей секции можно переписать как
Это первый момент, как меняется функционал, если мы не хотим предсказывать Σθ(xt,t), а фиксируем её.
Теперь вспомним, что q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I), но благодаря тому, что у нас гауссовское распределение, это можно переписать в виде
xt(x0,ϵ)=αˉtx0+1−αˉtϵ,ϵ∼N(0,I)
Выразим отсюда x0 и получим, что x0=αˉt1(xt−1−αˉtϵ), тогда подставим это выражение в формулу для μ~(xt,x0) (из подсекции «Вывод q(xt−1∣xt,x0)») и получим
μ~(xt,x0)=αt1(xt−1−αˉtβtϵ)
Теперь скажем, что наша модель будет предсказывать ϵ. И просто будем «подставлять» его в выражение для μ~ выше. Обозначим предсказание модели как ϵθ(xt,t) — предсказанный шум ϵ. Тогда лосс Lt превратиться в
Итак, алгоритмы обучения и сэмплирования выглядят вот так (на картинке z:=ϵ).
Алгоритм обучения и сэмплирования диффузионной модели (Изображение взято из: Ho et al. 2020)
Стоит отметить, что важным недостатком диффузионных моделей является низкая скорость сэмплирования. Согласно Song et al. 2020: «Требуется 20 часов на генерацию 50 тысяч картинок размера 32х32, используя DDPM, и меньше минуты, используя GAN» (Nvidia 2080 Ti GPU). Тем не менее, в данном направлении был достигнут значительный прогресс и в целом проблема медленного сэмплирования была частично решена: Jiaming Song et al. (2021), Kong & Ping (2021), Bond-Taylor et al. (2021)
Давайте зафиксируем, какие функции потерь можно использовать. Для всех них справедлив тот факт, что мы сэмплируем шаг равномерно во время обучение t∼U{1,...,T}) и оптимизируем соответствующий Lt.
Оптимизируя член из суммы LVLB. Это KL дивергенция между двумя нормальными распределениями
При фиксированной дисперсии и при предсказании шума с помощью взвешенной MSE. Или просто MSE. Ltsimple является самым популярным вариантом, который на практике дает лучшие результаты.
Расписание является гиперпараметром, основными требованиями на который являются невозрастание (β1≤...≤βT) и чтобы прямой процесс сходился к N(0,I) в пределе по T. Второе может гарантироваться тем, что αˉt→0. Вспомним,
q(xt∣x0)=N(xt;αˉtx0,(1−αˉt)I)
Однако на практике оно также проверяется, чтобы DKL(q(xT∣x0)∣∣N(0,I)) было близко к 0.
Также стоит упомянуть, что обычно берут T=1000. Но также важно помнить про требования выше, ведь расписание шума непосредственно зависит от T.
Чаще всего используют линейное расписание, где β1=10−4,βT=0.02. У данных констант нет никакой мотивации, кроме той, которая описана выше. Они были предложены в Ho et al. (2020).
В Nichol & Dhariwal (2021) было предложено косинусное расписание, которое помогло диффузионным моделям достичь лучшего NLL (negative loglikelihood):
Авторы обнаружили, что линейное расписание плохо работает на картинках 64х64 и меньше. А именно, последнии шаги прямого прохода были шумными и малоинформатиыными (просто зашумляем шум еще больше):
Пример зашумления картинки для линейного (сверху) и косинусного (снизу) расписания.
Также они обнаружили, что если обучать модель с линейным расписанием только на 80% первых шагов, то модель не становится сильно хуже, что подтверждает неиформативность последних шагов. Далее, они подобрали расписание так, чтобы αˉt убывало линейно на большей части отрезка (от 0 до T) и почти не менялось рядом с 0 и T. Разницу в αˉt для разных расписаний можно увидеть на картинке ниже:
Также они ограничивают βt числом 0.999, чтобы в конце процесса не было проблем с численной устойчивостью. Коэффициент s используется, чтобы βt не были слишком малы рядом с нулем. Он равен 0.008. Такое число было выбрано так, чтобы «β0 была немного меньше, чем размер бина одного пикселя, то есть 1/127.5»
Classifier guidance
В Nichol & Dhariwal (2021) был предложен метод условной генерации, который повышает качество генерируемых картинок, при этом уменьшая их разнообразие. Для этого предобучается «шумный» классификатор на зашумленных картинках, то есть pϕ(y∣xt). Затем он используется во время сэмплирования, корректируя предсказанное среднее на ∇xlogpϕ(y∣xt). В Nichol & Dhariwal (2021) (Секция 4.1) показывают, что данная добавка позволяет превратить распределение pθ(xi−1∣xi) в pθ(xi−1∣xi,y). Важно, что исходная диффузионная модель никак не меняется, что делает трюк еще более привлекательным. Алгоритм сэмплирования можно видеть на картинке ниже. Коэффициент s отвечает за силу guidance.
Мотивация
У генеративной модели GAN есть способ, который позволяет «балансировать» между разнообразием картинок и их качеством — truncation trick. Он заключается в сэмплировании латентного вектора truncated normal distibution. Данный трюк был хорошо описан и исследован в статье про BigGAN. Поэтому в диффузионных моделях тоже хотелось бы иметь метод, который позволяет балансировать между качеством и разнообразием. Авторы предложили classifier guidance, сравнили его с truncation trick и показали, что их метод строго лучше.
Ho & Salimans (2021) предложили метод, в котором guidance достигается без использования дополнительной модели, поскольку это достаточно затратно. Для этого они обучали условную модель ϵθ(xt∣y), у которой во время обучения реальная метка y заменялась с какой-то фиксированной вероятностью (10%) на пустую метку (y=∅). Это по сути позволяет нам обучать безусловную модель ϵθ(xt) одновременно с условной ϵθ(xt∣y)Тогда во время сэмплирования делаем так, чтобы предсказание немного менялось в сторону ϵθ(xt∣y), а именно:
ϵ^θ(xt∣y)=ϵθ(xt∣∅)+s⋅(ϵθ(xt∣y)−ϵθ(xt∣∅))
Мотивация этой формулы следовала из формулы Байеса:
Тогда мы можем просто подставить ∇xtlogp(y∣xt) в формулу для classifier guidance из предыдущей подсекции и получить желаемое равенство с точностью до коэффициента s.
Основная работа, в которой диффузионные модели (Denoising Diffusion Probabilistic Models, DDPMs) были применены для генерации картинок. Параграф в основном построен на ней.
Одна из первых попыток ускорить генерацию объектов. Идея следущая: давайте изменим прямой диффузионный процесс так, чтобы используя предобученную DDPM, приближать новый обратный процесс за меньшее число шагов.
Чтобы не обучать новую модель, нам нужен прямой диффузионный процесс, у которого будет такая же (суррогатная) функция потерь, а обратный процесс все еще останется Марковским. Оказалось, что существует целое семейство не-Марковских прямых процессов, удовлетворяющих этим требования. Это семейство имеет следующий вид:
Среднее было выбрано так, чтобы qσ(xt∣x0)=N(αtx0,(1−αt)I) для всех t. (см. Лемму 1 в Приложении B к статье). То есть важно лишь то, чтобы маргинальное распределение qσ(xt∣x0) не менялось по сравнению с обычным Марковским случаем. Прямой процесс может быть получен с помощью теоремы Байеса:
Тут σ контролирует степень стохастичности прямого процесса. Можно заметить, что в отличии от исходного диффузионного процесса, предложенный прямой процесс больше не является Марковским, так как каждый xt теперь зависит и от xt−1 и от x0. Схематично, это можно изобразить как на картинке справа. (Слева исходный диффузионный процесс для сравнения)
Заметка
Авторы обращают внимание, что функция потерь в DDPM зависит от q(xt∣x0), а не от q(x1:xT∣x0) напрямую. Это означает, что нам нужно выбрать любой другой прямой диффузионный процесс, у которого q(xt∣x0) остались те же.
Далее, мы можем переписать обратный процесс в данном виде:
Заметим, что при σt=(1−αt−1)(1−αt)1−αt/αt−1 прямой процесс становится марковским, а обратный как у DDPM (обычное сэмплирование, описанное в основной секции). При σt=0 процесс сэмплирования становится детерминистичным (данный способ и называется DDIM). Ускорение сэмплирования достигается засчет использования лишь какого-то подмножества шагов (0≤τ1≤...≤τS≤T,S<T). Также одним из плюсов детерминистичного сэмплирования является возможность делать семантическую интерполяцию в латентном пространстве (как у GANов).
Улучшение DDPM, в котором был предложен новое расписание шума, что улучшило NLL. Также был изучен вариант, в котором дисперсия Σθ(xt,t) предсказывается моделью.
Статья, в которой показывается, что DDPM могут генерировать более качественные картинки по сравнению с GANами. Также был предложен метод conditional сэмплирования. Для этого предобучается классификатор на зашумленных сэмплах, а во время сэмплирования среднее нормального распределения «корректируется» на градиент классификатора.
Диффузионные модели на дискретных данных (например, текст). Вместо нормальных распределений используются категориальные. Также была обобщена мультиномиальная диффузия с помощью «матриц перехода», которые задают способ зашумления дискретных данных.
Более подробно: у нас есть xt∈{1,...,K} — дискретная величина на всех шагах диффузии, тогда для каждого шага t определена матрица прямого перехода Qt такая, что [Qt]ij=q(xt=j∣xt−1=i). То есть строки матрицы суммируются в единицу. Тогда если обозначить через xt∈RK one-hot-закодированную версию xt, то прямой процесс можно описать через категориальные распределения:
Поскольку тут нет такой хорошей параметризации через ϵ, как у нормальных распределений, то единственный способ обучать — с помощью KL дивергенции (членамиLVLB).
Остается только понять, как выбирать Qt. Помимо того, чтобы сумма в каждой строчке была один, требуется, чтобы Qˉt сходилось (при t→∞) к равномерному распределению в каждой строчке (аналог нормального шума). За конкретными примерами стоит обратиться к статье.
Опишем работу метода GLIDE. Стоит задача генерировать картинки по заданному текстовому описанию. Для этого используется classifier-free guided diffusion model или CLIP. Это два разных варианта модели, которые авторы сравнивают. В первом случае модель обуславливается на эмбеддинги текста, которые были получены из обучаемого трансформера. Во втором случае guidance осуществляется за счет ∇xt⟨f(xt),g(c)⟩ (это по сути градиент лосса метода CLIP) . Тут f — это картиночный энкодер (на зашумленных картинках), а g — это энкодер текстового входа. В целом, авторы получили, что classifier-free guidance генерирует более качественные картинки.
